#!/usr/bin/env groovy

docker_registry_server = ImageTag.split(':')[0..1].join(':')
currentBuild.displayName = ImageTag.split(':')[2] + "-${arch}"
currentBuild.description = sourceBranch + ": " + commitHash

podTemplate(
cloud:'sc-ipp-blossom-prod',
yaml:"""
apiVersion: v1
kind: Pod
spec:
  volumes:
  - name: pvc-mount
    persistentVolumeClaim:
      claimName: 'kaolin-wisp-pvc'
  containers:
  - name: docker
    image: ${imageTag}
    command:
    - cat
    resources: 
        requests:
          nvidia.com/gpu: 1
        limits:
          nvidia.com/gpu: 1
    tty: true
    volumeMounts:
      - mountPath: /data
        name: pvc-mount
  imagePullSecrets:
  - name: gitlabregcred2
  nodeSelector:
    kubernetes.io/os: linux
    nvidia.com/gpu_type: ${arch}
""") {
  node(POD_LABEL) {
    container("docker") {
      updateGitlabCommitStatus(name: "test-${configName}-${arch}", state: 'running')
      def build_passed = true
      // Display some info
      try {
        timeout(time: 10, unit: 'SECONDS') {
            stage('Disp info') {
              sh 'nvidia-smi'
              sh 'python --version'
            }
        }
      } catch(e) {
        build_passed = false
        echo e.toString()
      }
      // Test NeRF with hashgrid
      try {
          timeout(time: 900, unit: 'SECONDS') {
              stage("NeRF-hashgrid lego (quick)") {
                  sh '''
                  cd /wisp
                  pytest -s \
                  --log-cli-level=INFO \
                  --log-cli-format="%(asctime)s|%(levelname)8s| %(message)s" \
                  tests/apps/test_nerf.py \
                  -k test_hashgrid_lego_quick \
                  --dataroot /data/ \
                  --dataset-num-workers 0
                  '''
              }
          }
      } catch(e) {
        build_passed = false
        echo e.toString()
      }
      try {
          timeout(time: 1200, unit: 'SECONDS') {
              stage("NeRF-hashgrid lego (best)") {
                  sh '''
                  cd /wisp
                  pytest -s \
                  --log-cli-level=INFO \
                  --log-cli-format="%(asctime)s|%(levelname)8s| %(message)s" \
                  tests/apps/test_nerf.py \
                  -k test_hashgrid_lego_best \
                  --dataroot /data/ \
                  --dataset-num-workers 0
                  '''
              }
          }
      } catch(e) {
        build_passed = false
        echo e.toString()
      }
      try {
          timeout(time: 1500, unit: 'SECONDS') {
              stage("NeRF-hashgrid V8") {
                  sh '''
                  cd /wisp
                  pytest -s \
                  --log-cli-level=INFO \
                  --log-cli-format="%(asctime)s|%(levelname)8s| %(message)s" \
                  tests/apps/test_nerf.py \
                  -k test_hashgrid_V8 \
                  --dataroot /data/ \
                  --dataset-num-workers 0
                  '''
              }
          }
      } catch(e) {
        build_passed = false
        echo e.toString()
      }
        // Check front facing hash-grid scene runs
        try {
        timeout(time: 260, unit: 'SECONDS') {
            stage("NeRF-hash standard-180") {
              sh '''
              cd /wisp
              WISP_HEADLESS=1 python app/nerf/main_nerf.py --config=app/nerf/configs/nerf_hash.yaml --dataset-path=/data/fox --max-epochs 5 --mip 2
              '''
            }
        }
        } catch(e) {
        build_passed = false
        echo e.toString()
        }

      // Train NeRF with octree
      try {
        timeout(time: 1500, unit: 'SECONDS') {
            stage("NeRF-octree lego") {
                  sh '''
                  cd /wisp
                  pytest -s \
                  --log-cli-level=INFO \
                  --log-cli-format="%(asctime)s|%(levelname)8s| %(message)s" \
                  tests/apps/test_nerf.py \
                  -k test_octree_lego \
                  --dataroot /data/ \
                  --dataset-num-workers 0
                  '''
            }
        }
      } catch(e) {
        build_passed = false
        echo e.toString()
      }
      // Train NeRF with octree
      try {
        timeout(time: 900, unit: 'SECONDS') {
            stage("NeRF-octree V8") {
                  sh '''
                  cd /wisp
                  pytest -s \
                  --log-cli-level=INFO \
                  --log-cli-format="%(asctime)s|%(levelname)8s| %(message)s" \
                  tests/apps/test_nerf.py \
                  -k test_octree_V8 \
                  --dataroot /data/ \
                  --dataset-num-workers 0
                  '''
            }
        }
      } catch(e) {
        build_passed = false
        echo e.toString()
      }

      // Train NeRF with triplanar feature maps
      try {
        timeout(time: 1500, unit: 'SECONDS') {
            stage("NeRF-triplanar lego") {
                  sh '''
                  cd /wisp
                  pytest -s \
                  --log-cli-level=INFO \
                  --log-cli-format="%(asctime)s|%(levelname)8s| %(message)s" \
                  tests/apps/test_nerf.py \
                  -k test_triplanar_lego \
                  --dataroot /data/ \
                  --dataset-num-workers 0
                  '''
            }
        }
      } catch(e) {
        build_passed = false
        echo e.toString()
      }
      // Train NeRF with VQAD
      try {
        timeout(time: 1500, unit: 'SECONDS') {
            stage("NeRF-codebook V8") {
                  sh '''
                  cd /wisp
                  pytest -s \
                  --log-cli-level=INFO \
                  --log-cli-format="%(asctime)s|%(levelname)8s| %(message)s" \
                  tests/apps/test_nerf.py \
                  -k test_codebook_V8 \
                  --dataroot /data/ \
                  --dataset-num-workers 0
                  '''
            }
        }
      } catch(e) {
        build_passed = false
        echo e.toString()
      }

      // Train NGLOD
      try {
        timeout(time: 350, unit: 'SECONDS') {
            stage("SDF-octree") {
              sh '''
              cd /wisp
              WISP_HEADLESS=1 python app/nglod/main_nglod.py --config=app/nglod/configs/nglod_octree.yaml --blas.mesh-path=/data/spot/spot_triangulated.obj
              '''
            }
        }
      } catch(e) {
        build_passed = false
        echo e.toString()
      }
      // Train SDFs with triplanar feature maps
      try {
        timeout(time: 750, unit: 'SECONDS') {
            stage("SDF-triplanar") {
              sh '''
              cd /wisp
              WISP_HEADLESS=1 python app/nglod/main_nglod.py --config=app/nglod/configs/nglod_triplanar.yaml --dataset.mesh-path=/data/spot/spot_triangulated.obj
              '''
            }
        }
      } catch(e) {
        build_passed = false
        echo e.toString()
      }
      // Train SDFs with hash grid
      try {
        timeout(time: 500, unit: 'SECONDS') {
            stage("SDF-hash") {
              sh '''
              cd /wisp
              WISP_HEADLESS=1 python app/nglod/main_nglod.py --config=app/nglod/configs/nglod_hash.yaml --blas.mesh-path=/data/spot/spot_triangulated.obj
              '''
            }
        }
      } catch(e) {
        build_passed = false
        echo e.toString()
      }

      // Example: latent_nerf
      try {
        timeout(time: 1500, unit: 'SECONDS') {
            stage("examples/latent_nerf") {
                  sh '''
                  cd /wisp
                  pytest -s \
                  --log-cli-level=INFO \
                  --log-cli-format="%(asctime)s|%(levelname)8s| %(message)s" \
                  tests/apps/test_latent_nerf.py \
                  --dataroot /data/ \
                  --dataset-num-workers 0
                  '''
            }
        }
      } catch(e) {
        build_passed = false
        echo e.toString()
      }

      if (build_passed) {
          updateGitlabCommitStatus(name: "test-${configName}-${arch}", state: 'success')
          currentBuild.result = "SUCCESS"
      } else {
          updateGitlabCommitStatus(name: "test-${configName}-${arch}", state: 'failed')
          currentBuild.result = "FAILURE"
      }
    }
  }
}
